#include "algo.h"
#include <cmath>
#include <fftw3.h>
#include <complex.h>
#include <string.h>

using namespace std;

static void multiply_fftw_complex(fftw_complex &p, fftw_complex q) {
	complex<double> c1, c2;
	c1.real(p[0]);
	c1.imag(p[1]);
	c2.real(q[0]);
	c2.imag(q[1]);
	c1 *= c2;

	p[0] = c1.real();
	p[1] = c1.imag();
}

vector<int> set_multiadd_limit_1d(vector<int> &S, vector<int> &T, int u) {
	int cnt, maxl;
	vector<int> res;
	fftw_complex *inS, *inT;
	fftw_plan plan;

	maxl = 2 * u + 1;

	inS = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * (maxl));
	memset(inS, 0, sizeof(fftw_complex) * maxl);

	inT = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * (maxl));
	memset(inT, 0, sizeof(fftw_complex) * maxl);

	for (cnt = 0; cnt < S.size(); ++cnt) {
		(inS[S[cnt]])[0] = 1;
	}

	for (cnt = 0; cnt < T.size(); ++cnt) {
		(inT[T[cnt]])[0] = 1;
	}

	/* Test code
 	for (cnt = 0; cnt < maxl; ++cnt) {
 		printf("%d: %F+%Fi\t", cnt, inS[cnt][0], inS[cnt][1]);
 		printf("%F+%Fi\n", inT[cnt][0], inT[cnt][1]);
 	}
 	putchar('\n');
 	*/

	plan = fftw_plan_dft_1d(maxl, inS, inS, FFTW_FORWARD, FFTW_ESTIMATE);
	fftw_execute(plan);
 	fftw_destroy_plan(plan);

	plan = fftw_plan_dft_1d(maxl, inT, inT, FFTW_FORWARD, FFTW_ESTIMATE);
	fftw_execute(plan);
 	fftw_destroy_plan(plan);

    /* Test code
 	for (cnt = 0; cnt < maxl; ++cnt) {
 		printf("%d: %F+%Fi\t", cnt, inS[cnt][0], inS[cnt][1]);
 		printf("%F+%Fi\n", inT[cnt][0], inT[cnt][1]);
 	}
 	putchar('\n');
 	*/

 	for (cnt = 0; cnt < maxl; ++cnt)
 		multiply_fftw_complex(inS[cnt], inT[cnt]);

    /*
 	for (cnt = 0; cnt < maxl; ++cnt) {
 		printf("%d: %F+%Fi\n", cnt, inS[cnt][0], inS[cnt][1]);
 	}
 	putchar('\n');
 	*/
 	

 	plan = fftw_plan_dft_1d(maxl, inS, inS, FFTW_BACKWARD, FFTW_ESTIMATE);
 	fftw_execute(plan);
 	fftw_destroy_plan(plan);

    /* Test code
 	for (cnt = 0; cnt < maxl; ++cnt) {
 		printf("%d: %F+%Fi\n", cnt, inS[cnt][0], inS[cnt][1]);
 	}
 	putchar('\n');
 	*/
    
 	for (cnt = 0; cnt <= u; ++cnt) {
 		if((int)(inS[cnt][0] / maxl + 0.5))
 			res.push_back(cnt);
 	}

    fftw_free(inS); 
    fftw_free(inT);
    return res;
}

vector<pair<int, int>> set_multiadd_limit_2d(vector<pair<int, int>> &S, vector<pair<int, int>> &T, int u, int n) {
	int i, j, maxl, maxm;
	vector<pair<int, int>> res;
	fftw_complex *inS, *inT;
	fftw_plan plan;

	maxl = 2 * u + 1;
	maxm = 2 * n + 1;

	inS = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * maxl * maxm);
	memset(inS, 0, sizeof(fftw_complex) * maxl *maxm);

	inT = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * maxl * maxm);
	memset(inT, 0, sizeof(fftw_complex) * maxl *maxm);

	for (i = 0; i < S.size(); ++i) {
		inS[S[i].first * maxm + S[i].second][0] = 1;
	}

	for (i = 0; i < T.size(); ++i) {
		inT[T[i].first * maxm + T[i].second][0] = 1;
	}

	plan = fftw_plan_dft_2d(maxl, maxm, inS, inS, FFTW_FORWARD, FFTW_ESTIMATE);
	fftw_execute(plan);
 	fftw_destroy_plan(plan);

	plan = fftw_plan_dft_2d(maxl, maxm, inT, inT, FFTW_FORWARD, FFTW_ESTIMATE);
	fftw_execute(plan);
 	fftw_destroy_plan(plan);

 	for (i = 0; i < maxl * maxm; ++i)
 		multiply_fftw_complex(inS[i], inT[i]);

 	plan = fftw_plan_dft_2d(maxl, maxm, inS, inS, FFTW_BACKWARD, FFTW_ESTIMATE);
 	fftw_execute(plan);
 	fftw_destroy_plan(plan);

    /* Test code
 	for (cnt = 0; cnt < maxl; ++cnt) {
 		printf("%d: %F+%Fi\n", cnt, inS[cnt][0], inS[cnt][1]);
 	}
 	putchar('\n');
 	*/
    
 	for (i = 0; i <= u; ++i)
 		for (j = 0; j <= n; ++j)
 		if((int)(inS[i * maxm + j][0] / maxl / maxm + 0.5))
 			res.push_back(pair<int, int>(i, j));

    fftw_free(inS); 
    fftw_free(inT);

 	return res;
}

vector<pair<int, int>> all_subsetsum_cardinality(vector<int> &nums, int begin, int end, int u) {
	vector<pair<int, int>> res;
	if (begin > end) {
		res.push_back(pair<int, int>(0, 0));
		return res;
	}

	if (begin == end) {
		res.push_back(pair<int, int>(0, 0));
		if (nums[begin] <= u)
			res.push_back(pair<int, int>(nums[begin], 1));
		return res;
	}

	int mid;
	vector<pair<int, int>> lf, rh;
	mid = (begin + end) / 2;

	lf = all_subsetsum_cardinality(nums, begin, mid, u);
	rh = all_subsetsum_cardinality(nums, mid + 1, end, u);

	res = set_multiadd_limit_2d(lf, rh, u, end - begin + 1);

	return res;
}

vector<int> all_subsetsum(vector<int> &nums, int u) {
	int cnt, cur;

	int b = sqrt(nums.size() * log2(nums.size()));
	vector<vector<int>> Q(b, vector<int>(0));
	vector<vector<pair<int, int>>> QC(b, vector<pair<int, int>>(0));
	vector<vector<int>> res(b, vector<int>(0));

	for (cnt = 0; cnt < nums.size(); ++cnt) {
		Q[nums[cnt] % b].push_back(nums[cnt] / b);
	}

	for (cnt = 0; cnt < b; ++cnt) {
		QC[cnt] = all_subsetsum_cardinality(Q[cnt], 0, Q[cnt].size() - 1, u / b);
		for (cur = 0; cur < QC[cnt].size(); ++cur) {
			res[cnt].push_back(QC[cnt][cur].first * b + QC[cnt][cur].second * cnt);

		}
	}

	for (cnt = 1; cnt < res.size(); ++cnt)
		res[0] = set_multiadd_limit_1d(res[0], res[cnt], u);

	return res[0];
}

int MaxSubSetSum(DataSet &s) {
	vector<int> res;
	res = all_subsetsum(s.nums, s.u);
	return res[res.size() - 1];
}
